/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.buffer;

import com.aduib.boot.common.concurrent.JctoolsHelper;
import com.aduib.boot.common.log.LOG;
import com.aduib.boot.common.util.SystemUtils;

import java.util.Collection;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.LongAdder;

import static com.aduib.boot.common.buffer.BufferChunk.isSubpage;
import static com.aduib.boot.common.buffer.PoolThreadCache.INTEGER_SIZE_MINUS_ONE;

/**
 * @description: BufferPoolArea
 * @author: zzh
 * @date: 2021/12/1 15:58
 */
public class BufferPoolArena implements BufferPool {

  static final int LOG2_QUANTUM = 4;
  private static final int LOG2_SIZE_CLASS_GROUP = 2;
  private static final int LOG2_MAX_LOOKUP_SIZE = 12;
  private static final int INDEX_IDX = 0;
  private static final int LOG2GROUP_IDX = 1;
  private static final int LOG2DELTA_IDX = 2;
  private static final int NDELTA_IDX = 3;
  private static final int PAGESIZE_IDX = 4;
  private static final int SUBPAGE_IDX = 5;
  private static final int LOG2_DELTA_LOOKUP_IDX = 6;
  private static final byte no = 0, yes = 1;
  private static final int FAIL_THRESHOLD = 1000;
  final int nSizes;
  // Number of thread caches backed by this arena.
  final AtomicInteger numThreadCaches = new AtomicInteger();

  final int directMemoryCacheAlignment;

  final int numSmallSubpagePools;

  private final short[][] sizeClasses;

  private final int[] pageIdx2sizeTab;

  // lookup table for sizeIdx <= smallMaxSizeIdx
  private final int[] sizeIdx2sizeTab;

  // lookup table used for size <= lookupMaxclass
  // spacing is 1 << LOG2_QUANTUM, so the size of array is lookupMaxclass >> LOG2_QUANTUM
  private final int[] size2idxTab;
  private final BufferChunkList q050;
  private final BufferChunkList q025;

  @Deprecated private final BufferChunkList q[];

  private final BufferChunkList q000;
  private final BufferChunkList qInit;
  private final BufferChunkList q075;
  private final BufferChunkList q100;
  private final BufferSubpage[] smallSubpagePools;
  private final int pageShifts;
  private final LongAdder allocationsSmall = new LongAdder();

  private final AtomicInteger chunkCount = new AtomicInteger(0);

  private final AtomicInteger failCount = new AtomicInteger(0);

  private final int pageSize;
  private final int chunkSize;
  private final AtomicLong capacity;
  private final AtomicLong size;
  private final ConcurrentMap<Thread, Integer> sharedOptsCount;
  /** 记录对线程ID->该线程的所使用Direct Buffer的size */
  private final ConcurrentMap<Long, Long> memoryUsage;
  private final int conReadBufferChunk;
  int nSubpages;
  int nPSizes;
  int smallMaxSizeIdx;
  private int lookupMaxSize;

  public BufferPoolArena(
      int chunkSize, int pageSize, int chunkCount, int conReadBufferChunk, int cacheAlignment) {

    this.pageSize = pageSize;
    this.pageShifts = chunkCount;
    this.chunkSize = chunkSize;

    int group = log2(chunkSize) + 1 - LOG2_QUANTUM;

    // generate size classes
    // [index, log2Group, log2Delta, nDelta, isMultiPageSize, isSubPage, log2DeltaLookup]
    sizeClasses = new short[group << LOG2_SIZE_CLASS_GROUP][7];
    nSizes = sizeClasses();

    // generate lookup table
    sizeIdx2sizeTab = new int[nSizes];
    pageIdx2sizeTab = new int[nPSizes];
    idx2SizeTab(sizeIdx2sizeTab, pageIdx2sizeTab);

    size2idxTab = new int[lookupMaxSize >> LOG2_QUANTUM];
    size2idxTab(size2idxTab);

    this.numSmallSubpagePools = nSubpages;
    this.directMemoryCacheAlignment = cacheAlignment;

    smallSubpagePools = newSubpagePoolArray(nSubpages);
    for (int i = 0; i < smallSubpagePools.length; i++) {
      smallSubpagePools[i] = newSubpagePoolHead();
    }

    q100 = new BufferChunkList(100, Integer.MAX_VALUE, chunkSize, this);
    q075 = new BufferChunkList(75, 100, chunkSize, this);
    q050 = new BufferChunkList(50, 100, chunkSize, this);
    q025 = new BufferChunkList(25, 75, chunkSize, this);
    q000 = new BufferChunkList(1, 50, chunkSize, this);
    qInit = new BufferChunkList(Integer.MIN_VALUE, 25, chunkSize, this);

    q100.nextList = null;
    q075.nextList = q100;
    q050.nextList = q075;
    q025.nextList = q050;
    q000.nextList = q025;
    qInit.nextList = q000;

    q100.prevList = q075;
    q075.prevList = q050;
    q050.prevList = q025;
    q025.prevList = q000;
    q000.prevList = null;
    qInit.prevList = qInit;

    this.chunkCount.set(chunkCount);
    this.conReadBufferChunk = conReadBufferChunk;

    /*--------------------------------------------------------------**/
    /** 废弃代码 */
    q = new BufferChunkList[6];
    q[5] = new BufferChunkList(100, Integer.MAX_VALUE, chunkSize, this);
    q[4] = new BufferChunkList(75, 100, chunkSize, this);
    q[3] = new BufferChunkList(50, 100, chunkSize, this);
    q[2] = new BufferChunkList(25, 75, chunkSize, this);
    q[1] = new BufferChunkList(1, 50, chunkSize, this);
    q[0] = new BufferChunkList(Integer.MIN_VALUE, 25, chunkSize, this);

    q[0].nextList = q[1];
    q[1].nextList = q[2];
    q[2].nextList = q[3];
    q[3].nextList = q[4];
    q[4].nextList = q[5];
    q[5].nextList = null;

    q[5].prevList = q[4];
    q[4].prevList = q[3];
    q[3].prevList = q[2];
    q[2].prevList = q[1];
    q[1].prevList = q[0];
    q[0].prevList = null;

    capacity = new AtomicLong(6L * chunkCount * chunkSize);
    size = new AtomicLong(6L * chunkCount * chunkSize);
    sharedOptsCount = JctoolsHelper.nonBlockingHashMap();
    memoryUsage = JctoolsHelper.nonBlockingHashMap();
    /*--------------------------------------------------------------**/
  }

  private static BufferPoolArena.bufferArea bufferArea(long handle) {
    return isSubpage(handle) ? bufferArea.Small : bufferArea.Normal;
  }

  private static Buffer allocateDirect(int capacity) {
    return SystemUtils.useDirectBufferNoCleaner()
        ? BufferUtils.allocateBuffer(SystemUtils.allocateDirectNoCleaner(capacity))
        : BufferUtils.allocateBuffer(capacity);
  }

  private BufferSubpage newSubpagePoolHead() {
    BufferSubpage head = new BufferSubpage();
    head.prev = head;
    head.next = head;
    return head;
  }

  @SuppressWarnings("unchecked")
  private BufferSubpage[] newSubpagePoolArray(int size) {
    return new BufferSubpage[size];
  }

  BufferSubpage findSubpageHead(int sizeIdx) {
    return smallSubpagePools[sizeIdx];
  }

  private int log2(int chunkSize) {
    return INTEGER_SIZE_MINUS_ONE - Integer.numberOfLeadingZeros(chunkSize);
  }

  private int sizeClasses() {
    int normalMaxSize = -1;

    int index = 0;
    int size = 0;

    int log2Group = LOG2_QUANTUM;
    int log2Delta = LOG2_QUANTUM;
    int ndeltaLimit = 1 << LOG2_SIZE_CLASS_GROUP;

    // First small group, nDelta start at 0.
    // first size class is 1 << LOG2_QUANTUM
    int nDelta = 0;
    while (nDelta < ndeltaLimit) {
      size = sizeClass(index++, log2Group, log2Delta, nDelta++);
    }
    log2Group += LOG2_SIZE_CLASS_GROUP;

    // All remaining groups, nDelta start at 1.
    while (size < chunkSize) {
      nDelta = 1;

      while (nDelta <= ndeltaLimit && size < chunkSize) {
        size = sizeClass(index++, log2Group, log2Delta, nDelta++);
        normalMaxSize = size;
      }

      log2Group++;
      log2Delta++;
    }

    // chunkSize must be normalMaxSize
    assert chunkSize == normalMaxSize;

    // return number of size index
    return index;
  }

  // calculate size class
  private int sizeClass(int index, int log2Group, int log2Delta, int nDelta) {
    short isMultiPageSize;
    if (log2Delta >= pageShifts) {
      isMultiPageSize = yes;
    } else {
      int pageSize = 1 << pageShifts;
      int size = (1 << log2Group) + (1 << log2Delta) * nDelta;

      isMultiPageSize = size == size / pageSize * pageSize ? yes : no;
    }

    int log2Ndelta = nDelta == 0 ? 0 : log2(nDelta);

    byte remove = 1 << log2Ndelta < nDelta ? yes : no;

    int log2Size = log2Delta + log2Ndelta == log2Group ? log2Group + 1 : log2Group;
    if (log2Size == log2Group) {
      remove = yes;
    }

    short isSubpage = log2Size < pageShifts + LOG2_SIZE_CLASS_GROUP ? yes : no;

    int log2DeltaLookup =
        log2Size < LOG2_MAX_LOOKUP_SIZE || log2Size == LOG2_MAX_LOOKUP_SIZE && remove == no
            ? log2Delta
            : no;

    short[] sz = {
      (short) index,
      (short) log2Group,
      (short) log2Delta,
      (short) nDelta,
      isMultiPageSize,
      isSubpage,
      (short) log2DeltaLookup
    };

    sizeClasses[index] = sz;
    int size = (1 << log2Group) + (nDelta << log2Delta);

    if (sz[PAGESIZE_IDX] == yes) {
      nPSizes++;
    }
    if (sz[SUBPAGE_IDX] == yes) {
      nSubpages++;
      smallMaxSizeIdx = index;
    }
    if (sz[LOG2_DELTA_LOOKUP_IDX] != no) {
      lookupMaxSize = size;
    }
    return size;
  }

  private void idx2SizeTab(int[] sizeIdx2sizeTab, int[] pageIdx2sizeTab) {
    int pageIdx = 0;

    for (int i = 0; i < nSizes; i++) {
      short[] sizeClass = sizeClasses[i];
      int log2Group = sizeClass[LOG2GROUP_IDX];
      int log2Delta = sizeClass[LOG2DELTA_IDX];
      int nDelta = sizeClass[NDELTA_IDX];

      int size = (1 << log2Group) + (nDelta << log2Delta);
      sizeIdx2sizeTab[i] = size;

      if (sizeClass[PAGESIZE_IDX] == yes) {
        pageIdx2sizeTab[pageIdx++] = size;
      }
    }
  }

  private void size2idxTab(int[] size2idxTab) {
    int idx = 0;
    int size = 0;

    for (int i = 0; size <= lookupMaxSize; i++) {
      int log2Delta = sizeClasses[i][LOG2DELTA_IDX];
      int times = 1 << log2Delta - LOG2_QUANTUM;

      while (size <= lookupMaxSize && times-- > 0) {
        size2idxTab[idx++] = i;
        size = idx + 1 << LOG2_QUANTUM;
      }
    }
  }

  private int alignSize(int size) {
    int delta = size & directMemoryCacheAlignment - 1;
    return delta == 0 ? size : size + directMemoryCacheAlignment - delta;
  }

  public int size2SizeIdx(int size) {
    if (size == 0) {
      return 0;
    }
    if (size > chunkSize) {
      return nSizes;
    }

    if (directMemoryCacheAlignment > 0) {
      size = alignSize(size);
    }

    if (size <= lookupMaxSize) {
      // size-1 / MIN_TINY
      return size2idxTab[size - 1 >> LOG2_QUANTUM];
    }

    int x = log2((size << 1) - 1);
    int shift =
        x < LOG2_SIZE_CLASS_GROUP + LOG2_QUANTUM + 1
            ? 0
            : x - (LOG2_SIZE_CLASS_GROUP + LOG2_QUANTUM);

    int group = shift << LOG2_SIZE_CLASS_GROUP;

    int log2Delta =
        x < LOG2_SIZE_CLASS_GROUP + LOG2_QUANTUM + 1 ? LOG2_QUANTUM : x - LOG2_SIZE_CLASS_GROUP - 1;

    int deltaInverseMask = -1 << log2Delta;
    int mod = (size - 1 & deltaInverseMask) >> log2Delta & (1 << LOG2_SIZE_CLASS_GROUP) - 1;

    return group + mod;
  }

  public int sizeIdx2size(int sizeIdx) {
    return sizeIdx2sizeTab[sizeIdx];
  }

  public int pages2pageIdx(int pages) {
    return pages2pageIdxCompute(pages, false);
  }

  public int pages2pageIdxFloor(int pages) {
    return pages2pageIdxCompute(pages, true);
  }

  private int pages2pageIdxCompute(int pages, boolean floor) {
    int pageSize = pages << pageShifts;
    if (pageSize > chunkSize) {
      return nPSizes;
    }

    int x = log2((pageSize << 1) - 1);

    int shift =
        x < LOG2_SIZE_CLASS_GROUP + pageShifts ? 0 : x - (LOG2_SIZE_CLASS_GROUP + pageShifts);

    int group = shift << LOG2_SIZE_CLASS_GROUP;

    int log2Delta =
        x < LOG2_SIZE_CLASS_GROUP + pageShifts + 1 ? pageShifts : x - LOG2_SIZE_CLASS_GROUP - 1;

    int deltaInverseMask = -1 << log2Delta;
    int mod = (pageSize - 1 & deltaInverseMask) >> log2Delta & (1 << LOG2_SIZE_CLASS_GROUP) - 1;

    int pageIdx = group + mod;

    if (floor && pageIdx2sizeTab[pageIdx] > pages << pageShifts) {
      pageIdx--;
    }

    return pageIdx;
  }

  @Override
  public Buffer allocate(PoolThreadCache cache, int reqCapacity, int maxCapacity) {
    PooledBuffer pooledBuffer = PooledBuffer.newInstance(maxCapacity);
    final int sizeIdx = size2SizeIdx(reqCapacity);
    if (sizeIdx <= smallMaxSizeIdx) {
      tcacheAllocateSmall(cache, pooledBuffer, reqCapacity, sizeIdx);
    } else if (sizeIdx < nSizes) {
      tcacheAllocateNormal(cache, pooledBuffer, reqCapacity, sizeIdx);
    } else {

    }
    //    tcacheAllocateNormal(cache, pooledBuffer, reqCapacity,sizeIdx);
    return pooledBuffer;
  }

  private void tcacheAllocateSmall(
      final PoolThreadCache cache,
      final PooledBuffer buf,
      final int reqCapacity,
      final int sizeIdx) {
    if (cache.allocateSmall(this, buf, reqCapacity, sizeIdx)) {
      // was able to allocate out of the cache so move on
      return;
    }
    final BufferSubpage head = smallSubpagePools[sizeIdx];
    final boolean needsNormalAllocation;
    synchronized (head) {
      final BufferSubpage s = head.next;
      needsNormalAllocation = s == head;
      if (!needsNormalAllocation) {
        assert s.doNotDestroy && s.elemSize == sizeIdx2size(sizeIdx);
        long handle = s.allocate();
        assert handle >= 0;
        s.chunk.initBufWithSubpage(buf, null, handle, reqCapacity, cache);
      }
    }

    if (needsNormalAllocation) {
      synchronized (this) {
        allocateNormal(buf, reqCapacity, cache, sizeIdx);
      }
    }

    incSmallAllocation();
  }

  private void tcacheAllocateNormal(
      final PoolThreadCache cache,
      final PooledBuffer buf,
      final int reqCapacity,
      final int sizeIdx) {
    if (cache.allocateNormal(this, buf, reqCapacity, sizeIdx)) {
      // was able to allocate out of the cache so move on
      return;
    }
    synchronized (this) {
      allocateNormal(buf, reqCapacity, cache, sizeIdx);
    }
  }

  private void allocateNormal(
      PooledBuffer buf, int reqCapacity, PoolThreadCache threadCache, int sizeIdx) {
    if (q050.allocate(buf, reqCapacity, sizeIdx, threadCache)
        || q025.allocate(buf, reqCapacity, sizeIdx, threadCache)
        || q000.allocate(buf, reqCapacity, sizeIdx, threadCache)
        || qInit.allocate(buf, reqCapacity, sizeIdx, threadCache)
        || q075.allocate(buf, reqCapacity, sizeIdx, threadCache)) {
      return;
    }

    BufferChunk chunk = newChunk(pageSize, nPSizes, pageShifts, chunkSize);
    chunk.allocate(buf, reqCapacity, sizeIdx, threadCache);
    qInit.add(chunk);
  }

  private void allocateHuge(PooledBuffer buf, int reqCapacity) {
    //    BufferChunk chunk = newUnpooledChunk(reqCapacity);
    //    activeBytesHuge.add(chunk.chunkSize);
    //    buf.initUnpooled(chunk, reqCapacity);
    //    allocationsHuge.increment();
  }

  private void incSmallAllocation() {
    allocationsSmall.increment();
  }

  @Override
  @Deprecated
  public Buffer allocate(int reqCapacity) {
    Buffer byteBuffer = null;
    int i = 0, count = 0;
    while (byteBuffer == null) {
      if (i > 5) {
        i = 0;
        count = failCount.incrementAndGet();
        if (count > FAIL_THRESHOLD) {
          expand();
        }
      }
      byteBuffer = q[i].allocate(reqCapacity);
      i++;
    }
    capacity.addAndGet(-reqCapacity);
    final Thread thread = Thread.currentThread();
    final long threadId = thread.getId();

    if (memoryUsage.containsKey(threadId)) {
      memoryUsage.put(threadId, memoryUsage.get(thread.getId()) + reqCapacity);
    } else {
      memoryUsage.put(threadId, (long) reqCapacity);
    }
    if (sharedOptsCount.containsKey(thread)) {
      int currentCount = sharedOptsCount.get(thread);
      currentCount++;
      sharedOptsCount.put(thread, currentCount);
    } else {
      sharedOptsCount.put(thread, 0);
    }
    return byteBuffer;
  }

  @Override
  public BigBuffer allocateBig(int reqCapacity) {
    return new BigBuffer(this, reqCapacity);
  }

  protected BufferChunk newChunk(int pageSize, int maxPageIdx, int pageShifts, int chunkSize) {
    if (directMemoryCacheAlignment == 0) {
      Buffer memory = allocateDirect(chunkSize);
      return new BufferChunk(pageSize, pageShifts, maxPageIdx, chunkSize, this, memory, memory);
    }

    final Buffer base = allocateDirect(chunkSize + directMemoryCacheAlignment);
    final Buffer memory =
        BufferUtils.allocateBuffer(
            SystemUtils.alignDirectBuffer(base.getByteBuffer(), directMemoryCacheAlignment));
    return new BufferChunk(pageSize, pageShifts, maxPageIdx, chunkSize, this, memory, base);
  }

  protected void destroyChunk(BufferChunk chunk) {
    if (SystemUtils.useDirectBufferNoCleaner()) {
      SystemUtils.freeDirectNoCleaner(chunk.memory.getByteBuffer());
    } else {
      SystemUtils.freeDirectBuffer(chunk.memory.getByteBuffer());
    }
  }

  @Override
  public BigBuffer allocateBig() {
    return new BigBuffer(this);
  }

  @Deprecated
  private void expand() {
    LOG.w("Current Buffer Size is not enough! Expanding Byte buffer!");
    BufferChunk byteBufferChunk =
        new BufferChunk(pageSize, chunkSize, this, BufferUtils.allocateBuffer(chunkSize));
    q[0].byteBufferChunks.add(byteBufferChunk);
    failCount.set(0);
  }

  @Override
  public void recycle(PoolThreadCache cache, Buffer buffer) {
    if (buffer instanceof PooledBuffer) {
      PooledBuffer localBuffer = (PooledBuffer) buffer;
      BufferChunk bufferChunk = localBuffer.chunk;
      bufferArea bufferArea = bufferArea(localBuffer.handle);
      if (cache != null) {
        if (cache.add(
            this, bufferChunk, localBuffer, localBuffer.handle, localBuffer.length, bufferArea)) {
          // cached so not free it.
          return;
        }
      }
      freeChunk(localBuffer.handle, localBuffer.maxLength, buffer, bufferChunk);
    } else {
      recycle(buffer);
    }
  }

  @Override
  @Deprecated
  public void recycle(Buffer byteBuffer) {
    final long size = byteBuffer != null ? byteBuffer.capacity() : 0;
    int i;
    for (i = 0; i < 6; i++) {
      if (q[i].free(byteBuffer)) {
        break;
      }
    }
    if (i > 5) {
      LOG.w("This ByteBuffer is not maintained in ByteBufferArena!");
      return;
    }
    final Thread thread = Thread.currentThread();
    final long threadId = thread.getId();

    if (memoryUsage.containsKey(threadId)) {
      memoryUsage.put(threadId, memoryUsage.get(thread.getId()) - size);
    }
    if (sharedOptsCount.containsKey(thread)) {
      int currentCount = sharedOptsCount.get(thread);
      currentCount--;
      sharedOptsCount.put(thread, currentCount);
    } else {
      sharedOptsCount.put(thread, 0);
    }
    capacity.addAndGet(byteBuffer.capacity());
  }

  public void free(PoolThreadCache cache, PooledBuffer buffer, long handle) {
    if (buffer instanceof PooledBuffer) {
      BufferChunk bufferChunk = buffer.chunk;
      bufferArea bufferArea = bufferArea(handle);
      if (cache != null) {
        if (cache.add(this, bufferChunk, buffer, handle, buffer.maxLength, bufferArea)) {
          // cached so not free it.
          return;
        }
      }
      freeChunk(handle, buffer.maxLength, buffer, bufferChunk);
    } else {
      recycle(buffer);
    }
  }

  /**
   * 释放chunk
   *
   * @param handle
   * @param maxLength
   * @param buffer
   * @param bufferChunk
   */
  public void freeChunk(long handle, int maxLength, Buffer buffer, BufferChunk bufferChunk) {
    boolean free = !bufferChunk.parent.free(handle, maxLength, bufferChunk, buffer);
    if (free) {
      destroyChunk(bufferChunk);
    }
  }

  @Override
  public long capacity() {
    return capacity.get();
  }

  @Override
  public long size() {
    return size.get();
  }

  @Override
  @Deprecated
  public int getConReadBufferChunk() {
    return conReadBufferChunk;
  }

  @Override
  @Deprecated
  public synchronized int getSharedOptsCount() {
    final Collection<Integer> integers = sharedOptsCount.values();
    int count = 0;
    for (int i : integers) {
      count += i;
    }
    return count;
  }

  /** @return */
  @Override
  public int getChunkSize() {
    return chunkSize;
  }

  @Override
  public int getPageSize() {
    return pageSize;
  }

  @Override
  @Deprecated
  public ConcurrentMap<Long, Long> getNetDirectMemoryUsage() {
    return memoryUsage;
  }

  enum bufferArea {
    Small,
    Normal
  }
}
